// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: donggu

#pragma once

#include <cstdlib>
#include <deque>
#include <functional>
#include <iostream>
#include <map>
#include <memory>
#include <set>
#include <string>
#include <utility>
#include <vector>

#include "coin2/base/config/app_config.h"
#include "coin2/base/log.h"
#include "coin2/base/time.h"
#include "coin2/exchange/base/api_base/raw_subscriber.h"
#include "coin2/exchange/base/executor/live_executor.h"
#include "coin2/exchange/base/feed/latency_manager.h"
#include "coin2/exchange/base/feed/ws/feed_subscription_spec.h"
#include "coin2/exchange/base/feed/ws/request_generator.h"
#include "coin2/exchange/base/feed/ws/ws_subscriber.h"
#include "coin2/exchange/base/feed/ws/rest_subscriber.h"
#include "coin2/exchange/base/feed/ws/fix_subscriber.h"
#include "coin2/exchange/base/feed/ws/kucoin_subscriber.h"
#include "coin2/exchange/base/feed/ws/bittrex_subscriber.h"
#include "coin2/exchange/base/market/market_exchange_api.h"
#include "coin2/exchange/factory/feed/request_generator_factory.h"
#include "coin2/exchange/feed.pb.h"
#include "coin2/exchange/util/rate_limit.h"
#include "coin2/flow/symbol_group.h"
#include "coin2/flow/topic_string.h"

namespace coin2::exchange::base::feed {

class SubscriberManager : public std::enable_shared_from_this<SubscriberManager>,
                     public api_base::BaseRawAsyncSubscriber {
 public:
  SubscriberManager(
      const FeedSubscriptionSpec& spec,
      const WebsocketConfig& config,
      coin2::flow::TopicType topic_type,
      api_base::LiveAsyncExecutor* flow_sub,
      bool enable_telemetry)
      : BaseRawAsyncSubscriber(flow_sub),
        spec_(spec),
        config_(config),
        topic_type_(topic_type),
        enable_telemetry_(enable_telemetry) {
  }


  void RegisterPubTopic() override {
    for (auto& subsriber : subscribers_) {
      subsriber->RegisterPubTopic();
    }
  }

  void RegisterPubTopic(const std::string& topic) override {
    for (auto& subsriber : subscribers_) {
      subsriber->RegisterPubTopic(topic);
    }
  }

  void SubscribeAsync() override {
    if (connection_configs_.empty()) {
      CreateSubscribers();
    }
    StartSubscribersAsync();
  };

  void GetHealthReport(int indent, HealthReport* health) const override {
    for (auto& subsriber : subscribers_) {
      subsriber->GetHealthReport(indent, health);
    }
  }

  void CleanUp() override {
    for (auto& subsriber : subscribers_) {
      subsriber->CleanUp();
    }
  }

  std::vector<WsConnectionConfig> GenerateConnectionConfigs() {
    std::vector<WsConnectionConfig> connection_configs;
    BuildWsConnectionConfig(config_, spec_, native_symbol_groups_,
                            topic_type_, &connection_configs);
    return connection_configs;
  }

  bool BanProduct(const std::vector<const IProduct*>& products) override {
    // remove from local config
    for (const auto& prod : products) {
      auto symbol = prod->feed_native_symbol();
      char group = coin2::flow::SymbolGroupConfig::GetInstance().GetGroupForRelativeNorm(
          spec_.mea,
          prod->relative_norm());
      RemoveNativeSymbol(symbol, group);
    }

    // remove from children
    for (auto& subsriber : subscribers_) {
      subsriber->BanProduct(products);
    }

    return true;
  }

 private:
  static std::string BuildTopic(
    const FeedSubscriptionSpec& spec,
    coin2::flow::TopicType topic_type,
    char group,
    int worker,
    int connection_id) {
    std::string topic;
    if (topic_type == coin2::flow::TopicType::LONG_TOPIC) {
      topic = spec.mea.SnakeString() + "_" + group + std::to_string(worker);
    } else if (topic_type == coin2::flow::TopicType::SHORT_TOPIC) {
      topic = spec.mea.SnakeString();
    } else if (topic_type == coin2::flow::TopicType::CONNECTION_TOPIC) {
      topic = spec.mea.SnakeString() + "_" + group + "_" +
              std::to_string(connection_id) + "_" + std::to_string(worker);
    } else {
      CHECK(false) << "Invalid topic_type: " << topic_type;
    }

    return topic;
  }

  static std::vector<std::string> channel_intersection(std::set<std::string>* v1,
                                                std::set<std::string>* v2) {
    std::vector<std::string> v3;
    std::set_intersection(v1->begin(), v1->end(),
                          v2->begin(), v2->end(),
                          back_inserter(v3));
    return v3;
  }

  static void BuildWsConnectionConfig(
      const WebsocketConfig& config,
      const FeedSubscriptionSpec& spec,
      const std::map<char, std::vector<std::string>>& native_symbol_groups,
      coin2::flow::TopicType topic_type,
      std::vector<WsConnectionConfig>* configs) {
    std::vector<WsUrl> url_configs;
    std::set<std::string> channels;
    for (auto& ch : spec.channels) {
      channels.insert(ch.first);
    }
    // build by url
    for (auto& url : spec.extra_urls) {
      std::string url_name = url.first;
      WsUrl url_config = url.second;
      // config only override websocket url
      if (url_config.host.find("https://") == std::string::npos) {
        if (config.has_ws_host()) url_config.host = config.ws_host();
        size_t first_slash = url_config.host.find(R"(/)");
        if (first_slash != std::string::npos && url_config.host.find("fix://") == std::string::npos) {
          url_config.path = url_config.host.substr(first_slash) + url_config.path;
          url_config.host = url_config.host.substr(0, first_slash);
        }
        if (config.has_ws_host_port()) url_config.port = config.ws_host_port();
      }
      std::set<std::string> url_channels(
          url_config.channels.begin(),
          url_config.channels.end());
      std::set<std::string> config_channels(
          config.channels().begin(),
          config.channels().end());

      url_config.channels = channel_intersection(&url_channels, &config_channels);
      for (auto& ch : url_channels) {
        channels.erase(ch);
      }

      auto& urls = config.urls();
      auto iter = urls.find(url_name);
      if (iter != urls.end()) {
        if (iter->second.has_unauthorized_limit()) {
          url_config.unauthorized_limit.subscription =
              RateFromIso8601(iter->second.unauthorized_limit().subscription_rate());
        }
        if (iter->second.has_rest_host()) {
          url_config.host = iter->second.rest_host();
          size_t first_slash = url_config.host.find(R"(/)", 8);
          if (first_slash != std::string::npos) {
            url_config.path = url_config.host.substr(first_slash) + url_config.path;
            url_config.host = url_config.host.substr(0, first_slash);
          }
        }
        if (iter->second.has_ws_host()) {
          url_config = url.second;
          url_config.host = iter->second.ws_host();
          size_t first_slash = url_config.host.find(R"(/)");
          if (first_slash != std::string::npos) {
            url_config.path = url_config.host.substr(first_slash) + url_config.path;
            url_config.host = url_config.host.substr(0, first_slash);
          }
        }
        if (iter->second.has_http_proxy()) {
          url_config.proxy = iter->second.http_proxy();
        }
      }

      if (!url_config.channels.empty()) {
        url_configs.push_back(url_config);
      }
    }

    // default url setting
    WsUrl url_config = {spec.host, spec.port, spec.path, {}, {} /* proxy */,
        "default", {}, {}, {}, false,
        {}, false, spec.support_ping_frame};
    if (config.has_ws_host()) url_config.host = config.ws_host();
    size_t first_slash = url_config.host.find(R"(/)");
    if (first_slash != std::string::npos && url_config.host.find("fix://") == std::string::npos) {
      url_config.path = url_config.host.substr(first_slash) + url_config.path;
      url_config.host = url_config.host.substr(0, first_slash);
    }
    if (config.has_ws_host_port()) url_config.port = config.ws_host_port();
    if (config.has_http_proxy()) url_config.proxy = config.http_proxy();
    std::set<std::string> config_channels(
        config.channels().begin(),
          config.channels().end());
    url_config.channels = channel_intersection(&channels, &config_channels);
    if (!url_config.channels.empty()) {
      url_configs.push_back(url_config);
    }

    // build by group
    int connection_id = 0;
    bool disable_ssl = config.has_ws_disable_ssl()? config.ws_disable_ssl() : false;
    std::vector<WsConnectionConfig> connection_configs;
    for (auto& url : url_configs) {
      for (auto& it : native_symbol_groups) {
        int symbols_per_connection = spec.symbols_per_connection;
        if (config.has_symbols_per_connection()) {
          symbols_per_connection = config.symbols_per_connection();
        }
        bool require_auth = spec.require_auth;
        if (config.has_require_auth()) {
          // override
          require_auth = config.require_auth();
        }
        bool require_header = spec.require_header;
        if (config.has_require_header()) {
          // override
          require_header = config.require_header();
        }
        if (symbols_per_connection > 0) {
          for (int i = 0; i < it.second.size(); i += symbols_per_connection) {
            WsConnectionConfig con_config;
            con_config.connection_id = connection_id;
            connection_id++;
            con_config.url = url;
            con_config.url.disable_ssl = disable_ssl;
            for (int j = 0; j < symbols_per_connection && (i + j) < it.second.size(); ++j) {
              con_config.symbols.push_back(it.second[i + j]);
            }
            con_config.require_auth = require_auth;
            con_config.require_header = require_header;
            con_config.group = it.first;
            con_config.mea_config = &spec;
            con_config.ws_config = &config;
            connection_configs.push_back(con_config);
          }
        } else {
          WsConnectionConfig con_config;
          con_config.connection_id = connection_id;
          connection_id++;
          con_config.url = url;
          con_config.url.disable_ssl = disable_ssl;
          con_config.symbols = it.second;
          con_config.require_auth = require_auth;
          con_config.require_header = require_header;
          con_config.group = it.first;
          con_config.mea_config = &spec;
          con_config.ws_config = &config;
          connection_configs.push_back(con_config);
        }
      }
    }

    // build by worker
    int start_worker_id = 1;
    if (config.has_worker_id()) start_worker_id = config.worker_id();
    int end_worker_id = start_worker_id + config.num_workers();
    for (int w = start_worker_id; w < end_worker_id; w++) {
      int count = 0;
      for (auto& con_config : connection_configs) {
        std::string topic = BuildTopic(*con_config.mea_config, topic_type, con_config.group, w, con_config.connection_id);
        // Prex case always use long topic name.
        if (con_config.mea_config->mea.exchange == coin::proto::ExchangeType::Prex) {
          topic = BuildTopic(
              *con_config.mea_config, coin2::flow::TopicType::LONG_TOPIC, con_config.group, w, con_config.connection_id);
        }
        WsConnectionConfig connection_config = con_config;
        connection_config.topic = topic;
        connection_config.worker_id = w;
        if (config.bind_ips_size() >= config.num_workers()) {
          int half_ips_count = config.bind_ips_size() / 2;
          connection_config.url.bind_ip = config.bind_ips(
              w - 1 + (count % half_ips_count * half_ips_count));
        }
        if (config.has_bind_ip()) {
          connection_config.url.bind_ip = config.bind_ip();
        }
        if (config.key_filepaths_size() >= config.num_workers() &&
            config.key_filepaths_size() > (w - 1)) {
          connection_config.key_filepath = config.key_filepaths(w - 1);
        } else if (config.key_filepaths_size() > 0) {
          connection_config.key_filepath = config.key_filepaths(0);
        }

        if (config.ws_hosts_size() > (w - 1)) {
          connection_config.url.host = config.ws_hosts(w - 1);
        }
        configs->push_back(connection_config);
        ++count;
      }
    }
  }

  static void CopyNativeSymbol(
      const std::map<char, std::vector<std::string>> & symbol_groups,
      BaseRawAsyncSubscriber* subscriber){
    for (auto const& [group, symbols] : symbol_groups) {
      for (auto const& symbol : symbols) {
        subscriber->RequestNativeSymbol(symbol, group);
      }
    }
  }

  void StartSubscribersAsync() {
    for (auto& subsriber : subscribers_) {
      subsriber->StartAsync();
    }
  }

 public:
  void CreateSubscribers() {
    BuildWsConnectionConfig(config_, spec_, native_symbol_groups_, topic_type_,
                          &connection_configs_);

    std::vector<WsConnectionConfig> fix_config;
    std::map<std::string, std::vector<WsConnectionConfig>> rest_config;
    std::vector<WsConnectionConfig> websocket_config;
    std::vector<WsConnectionConfig> kucoin_config;
    std::vector<WsConnectionConfig> bittrex_config;

    // should classify by url type
    // websocket use 1 subscriber, rest api use 1 subscriber
    for (const WsConnectionConfig& config : connection_configs_) {
      if (config.url.host.rfind("fix://") == 0) {
          // fix api starts with "fix://"
          fix_config.emplace_back(config);
      } else if (config.url.host.rfind("https://") == 0) {
          // rest api starts with "https://"
          rest_config[config.url.name].emplace_back(config);
      } else if (config.mea_config->mea.exchange == coin::proto::ExchangeType::Kucoin) {
        kucoin_config.emplace_back(config);
      } else if (config.mea_config->mea.exchange == coin::proto::ExchangeType::Bittrex) {
        bittrex_config.emplace_back(config);
      } else {
        // default is websocket
        websocket_config.emplace_back(config);
      }
    }

    if (!fix_config.empty()) {
      std::shared_ptr<FixSubscriber> fix_subscriber = std::make_shared<FixSubscriber>(
          spec_,
          fix_config,
          flow_sub_);
      CopyNativeSymbol(native_symbol_groups(), fix_subscriber.get());
      subscribers_.push_back(std::static_pointer_cast<BaseRawAsyncSubscriber>(fix_subscriber));
    }

    if (!rest_config.empty()) {
      for (auto &config : rest_config) {
        std::shared_ptr<RestSubscriber> rest_subcriber = std::make_shared<RestSubscriber>(
            spec_,
            config_,
            config.second,
            flow_sub_);

        CopyNativeSymbol(native_symbol_groups(), rest_subcriber.get());
        subscribers_.push_back(std::static_pointer_cast<BaseRawAsyncSubscriber>(rest_subcriber));
      }
    }

    if (!websocket_config.empty()) {
      std::shared_ptr<WsSubscriber> ws_subcriber = std::make_shared<WsSubscriber>(
          spec_,
          config_,
          websocket_config,
          flow_sub_,
          enable_telemetry_);
      CopyNativeSymbol(native_symbol_groups(), ws_subcriber.get());
      subscribers_.push_back(std::static_pointer_cast<BaseRawAsyncSubscriber>(ws_subcriber));
    }

    if (!kucoin_config.empty()) {
      std::shared_ptr<KucoinSubscriber> kucoin_subcriber = std::make_shared<KucoinSubscriber>(
          spec_,
          config_,
          kucoin_config,
          flow_sub_,
          enable_telemetry_);
      CopyNativeSymbol(native_symbol_groups(), kucoin_subcriber.get());
      subscribers_.push_back(std::static_pointer_cast<BaseRawAsyncSubscriber>(kucoin_subcriber));
    }

    if (!bittrex_config.empty()) {
      std::shared_ptr<BittrexSubscriber> bittrex_subcriber = std::make_shared<BittrexSubscriber>(
          spec_,
          config_,
          bittrex_config,
          flow_sub_,
          enable_telemetry_);
      CopyNativeSymbol(native_symbol_groups(), bittrex_subcriber.get());
      subscribers_.push_back(std::static_pointer_cast<BaseRawAsyncSubscriber>(bittrex_subcriber));
    }
  }

  const std::set<std::string> GetTopics() const {
    std::set<std::string> topics;
    for (const auto& conn : connection_configs_) {
      topics.insert(conn.topic);
    }
    return topics;
  }

 private:
  const FeedSubscriptionSpec& spec_;
  const WebsocketConfig& config_;
  coin2::flow::TopicType topic_type_;
  bool enable_telemetry_;
  std::vector<std::shared_ptr<BaseRawAsyncSubscriber>> subscribers_;
  std::vector<WsConnectionConfig> connection_configs_;
};

}  // namespace coin2::exchange::base::feed
